#
# Functions for working with words
#

from collections import defaultdict, OrderedDict
import string
import pysrt
from bs4 import BeautifulSoup
from stemming.porter2 import stem as get_stem

MIN_LEN = 3
MAX_LEN = 30
PUNCT_CHARS = ',.:;!?"\''


def get_words_count_from_srt(file_name):
    subs = pysrt.open(file_name, encoding='iso-8859-1')
    subs_raw_text = " ".join([sub.text for sub in subs])

    soup = BeautifulSoup(subs_raw_text)
    subs_plain_text = soup.get_text()    

    words_list = map(lambda w: w.lower().strip(PUNCT_CHARS), subs_plain_text.split())
    words_list = filter(lambda w: MIN_LEN < len(w) < MAX_LEN, words_list)

    words_count = defaultdict(int)
    for word in words_list:
        words_count[word] += 1

    return words_count


def stem_words_count(words_count):
    stem_groups = {}
    for word, count in words_count.items():
        stem = get_stem(word)
        if stem not in stem_groups:
            stem_groups[stem] = {
                'words': set(),
                'count': 0
            }
        stem_groups[stem]['words'].add(word)
        stem_groups[stem]['count'] += count

    result = {}
    for stem, group in stem_groups.items():
        shortest_form = sorted(group['words'], key=len)[0]
        result[shortest_form] = group['count']

    return result


def get_most_frequent_words(words_count, n=100):
    if n == -1:
        n = len(words_count)

    return OrderedDict((sorted(words_count.items(), key=lambda x: x[1], reverse=True))[0:n])

